#!/usr/bin/python3
# -*- coding: utf-8 -*-

import sys
# sys.path.append('/opt/work/caffe/python')
sys.path.insert(0, '.')

import os
import numpy as np
import random
import json


# base_path = r'/rootfs/media/yery/Kaso/data'
base_path = r'/rootfs/media/kasim/Data1/data'

image_base_path = os.path.join(base_path, 'ErisedPersonGender')
json_dir_path = os.path.join(base_path, 'ErisedPersonGender/Annotations.20201105')
# json_dir_path = os.path.join(base_path, 'ErisedPersonGender/Annotations.20201029')
# file_faceid_list_path = os.path.join(base_path, 'ErisedPersonGender/file_faceid_list.txt')
file_faceid_list_path = os.path.join(base_path, 'VideoGenderCropPerson/FaceID/Mtypical_file_faceid_list.txt')
out_path = os.path.join(base_path, 'VideoGenderCropPerson/FaceID')

age_group_map = {
    '0-2岁': 0,  # AgeGroup.UNDER_TWO,
    '3-5岁': 1,  # AgeGroup.THREE_TO_FIVE,
    '6-12岁': 2,  # AgeGroup.SIX_TO_TWELVE,
    '13-18岁': 3,  # AgeGroup.THIRTEEN_TO_EIGHTEEN,
    '19-23岁': 4,  # AgeGroup.NINETEEN_TO_TWENTY_THREE,
    '24-29岁': 5,  # AgeGroup.TWENTY_FOUR_TO_TWENTY_NINE,
    '30-35岁': 6,  # AgeGroup.THIRTY_TO_THIRTY_FIVE,
    '36-45岁': 7,  # AgeGroup.THIRTY_SIX_TO_FORTY_FIVE,
    '46-60岁': 8,  # AgeGroup.FORTY_SIX_TO_SIXTY,
    '60岁以上': 9,  # AgeGroup.ABOVE_SIXTY,
    # '跳过': None,  # AgeGroup.UNCERTAIN,
}


def main():
    random.seed(123456)
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    faceid_files = {}
    with open(file_faceid_list_path, 'r') as file:
        for line in file.readlines():
            # gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4be6600/video/20190926_085306/00043_000.jpg 1132
            lines = line.strip().split()
            file_name = lines[0]
            faceid = int(lines[1])
            file_names = faceid_files.get(faceid, None)
            if file_names is None:
                faceid_files[faceid] = [file_name]
            else:
                file_names.append(file_name)

    print('FaceID Count:', len(faceid_files.keys()))

    choice_count = 0
    choice_faceids = []
    all_choice_faceids = []
    one_faceids = []
    age_group_count = len(age_group_map.keys())
    total_file_count = 0
    filename_faceid_agegroups = []
    for faceid, file_names in faceid_files.items():
        file_count = len(file_names)
        total_file_count += file_count
        age_group_counts = [0] * age_group_count
        for file_name in file_names:
            json_file_name = os.path.splitext(file_name.replace('/', '_'))[0] + '.json'
            json_file_path = os.path.join(json_dir_path, json_file_name)
            with open(json_file_path, 'r', encoding='utf-8') as f:
                json_info = json.load(f)
                if json_info is None:
                    continue
                name = json_info['outputs']['object'][0]['name']
                age_group = age_group_map.get(name, None)
                if age_group is None:
                    continue
            age_group_counts[age_group] += 1
        age_group_counts = np.array(age_group_counts, dtype=np.int32)
        age_group_counts_mask = age_group_counts > 0
        age_group_counts_count = age_group_counts_mask.sum()
        if age_group_counts_count == 0:
            age_group = 10
        elif age_group_counts_count == 1:
            age_group = np.nonzero(age_group_counts_mask)[0][0]
            one_faceids.append([faceid, age_group])
        else:
            age_group_count_idxes = np.argsort(-age_group_counts)
            age_group0 = age_group_count_idxes[0]
            age_group1 = age_group_count_idxes[1]
            if age_group_counts[age_group0] > age_group_counts[age_group1]:
                age_group = age_group0
            elif age_group_counts[age_group0] == age_group_counts[age_group1]:
                if age_group_counts_count == 2:
                    age_group = random.choice([age_group0, age_group1])
                    choice_count += 1
                    choice_faceids.append([faceid, age_group])
                else:
                    age_groups = np.nonzero(age_group_counts_mask)[0]
                    age_group_start = age_groups[0]
                    age_group_end = age_groups[-1]
                    age_group_mid = (age_group_start + age_group_end) / 2
                    age_group0_dif = abs(age_group0 - age_group_mid)
                    age_group1_dif = abs(age_group1 - age_group_mid)
                    if age_group0_dif == age_group1_dif:
                        age_group = random.choice([age_group0, age_group1])
                        choice_count += 1
                        choice_faceids.append([faceid, age_group])
                    elif age_group0_dif < age_group1_dif:
                        age_group = age_group0
                    else:
                        age_group = age_group1
            else:
                print('Error:', faceid, file_names)
                continue
            all_choice_faceids.append([faceid, age_group])
        for file_name in file_names:
            filename_faceid_agegroups.append([file_name, faceid, int(age_group)])

    print('Choice Count:', choice_count)

    filename_faceid_agegroups.sort(key=lambda x: x[0])
    out_file_path = os.path.join(out_path, 'file_faceid_agegroup_list.txt')
    with open(out_file_path, 'w') as file:
        for filename, faceid, agegroups in filename_faceid_agegroups:
            out_info = '{} {} {}\n'.format(filename, faceid, agegroups)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    choice_faceids.sort(key=lambda x: x[0])
    out_file_path = os.path.join(out_path, 'choice_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for faceid, age_group in choice_faceids:
            out_info = '{} {}\n'.format(faceid, age_group)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    all_choice_faceids.sort(key=lambda x: x[0])
    out_file_path = os.path.join(out_path, 'all_choice_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for faceid, age_group in all_choice_faceids:
            out_info = '{} {}\n'.format(faceid, age_group)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    one_faceids.sort(key=lambda x: x[0])
    out_file_path = os.path.join(out_path, 'one_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for faceid, age_group in one_faceids:
            out_info = '{} {}\n'.format(faceid, age_group)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    print('Finish!')


if __name__ == '__main__':
    main()
